from sqlalchemy import create_engine, text
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, scoped_session
import os

# 创建基类，供模型继承
Base = declarative_base()

class Database:
    def __init__(self, db_name='sqlite', root_path=None):
        """
        初始化数据库连接
        :param db_name: 数据库名称 (默认为 'sqlite')
        :param root_path: 数据库文件存储根路径
        """
        if root_path is None:
            root_path = os.getcwd()
            
        self.engine = self._create_engine(db_name, root_path)
        self.Session = sessionmaker(bind=self.engine)
        self.session = None
        
    def _create_engine(self, db_name, root_path):
        # SQLite 连接字符串
        db_path = os.path.join(root_path, f"{db_name}.db")  # 使用 .db 后缀
        return create_engine(f"sqlite:///{db_path}")
    
    def init_app(self, app):
        """
        用于Flask应用初始化
        """
        # 确保表存在
        Base.metadata.create_all(self.engine)
        
        # 设置应用关闭时清理会话
        @app.teardown_appcontext
        def shutdown_session(exception=None):
            if self.session:
                self.session.close()
    
    def get_session(self):
        """获取数据库会话"""
        if not self.session:
            self.session = scoped_session(self.Session)
        return self.session
    
    def insert(self, model_instance):
        """插入数据"""
        session = self.get_session()
        try:
            session.add(model_instance)
            session.commit()
            session.refresh(model_instance)
            return model_instance
        except Exception as e:
            session.rollback()
            raise e
    
    def select(self, model_class, **filters):
        """查询数据"""
        session = self.get_session()
        try:
            query = session.query(model_class)
            for key, value in filters.items():
                query = query.filter(getattr(model_class, key) == value)
            return query.all()
        except Exception as e:
            session.rollback()
            raise e
    
    def update(self, model_class, updates, **filters):
        """更新数据"""
        session = self.get_session()
        try:
            query = session.query(model_class)
            for key, value in filters.items():
                query = query.filter(getattr(model_class, key) == value)
            query.update(updates)
            session.commit()
        except Exception as e:
            session.rollback()
            raise e
    
    def delete(self, model_class, **filters):
        """删除数据"""
        session = self.get_session()
        try:
            query = session.query(model_class)
            for key, value in filters.items():
                query = query.filter(getattr(model_class, key) == value)
            query.delete()
            session.commit()
        except Exception as e:
            session.rollback()
            raise e
    
    def execute_sql(self, sql):
        """执行原始SQL"""
        session = self.get_session()
        try:
            result = session.execute(text(sql))
            session.commit()
            return result.fetchall()
        except Exception as e:
            session.rollback()
            raise e
        finally:
            session.close()

# 创建默认数据库实例
db = Database() 